# -*- coding: UTF-8 -*-
# Author: Liang Bing
# E-mail: believer19940901@gmail.com
# Date Time: 2025/6/29 19:18

import json
from pathlib import Path

import click
import numpy as np
import pandas as pd

from biobee.configs import RSCRIPT, SCRIPTS
from biobee.utils import run_or_die, open_excel


def func(x, p_value, fc):
    if x["PValue"] >= p_value:
        return "not"
    if x["logFC"] > np.log2(fc):
        return "up"
    if x["logFC"] < -np.log2(fc):
        return "down"
    return "not"


def mark_top(group, top):
    group["Regular"] = group.name
    if group.name == 'not':
        # 分组名是 'not'，全部标记为 'F'
        group['flag'] = 'F'
    else:
        # 正常分组，标记 logFC_abs 最大为 'T'，其他为 'F'
        top = group.nlargest(top, 'abs_logFC')
        group['flag'] = group['Gene'].isin(top['Gene']).map({True: 'T', False: 'F'})
    return group


def convert_input(input_path: Path, output_path: Path, p_value, fc, top):
    df: pd.DataFrame = open_excel(input_path)
    df.columns = ["Gene", "logFC", "PValue"]
    df["Regular"] = df.apply(func, args=(p_value, fc, ), axis=1)
    df["abs_logFC"] = df.logFC.abs()
    df["Rank"] = df["logFC"].rank(method="first", ascending=False).astype(int)
    df = df.groupby('Regular', as_index=False).apply(mark_top, top, include_groups=False)
    print(df["flag"].value_counts())
    df["neg_log10PValue"] = -np.log10(df["PValue"])
    df.to_csv(output_path, index=False, header=True, encoding="utf-8")
    click.echo(f"[✓] 输入文件已转换为 CSV：{output_path}")


@click.command("DiffRank", help="差异排序图")
def diff_rank():
    try:
        config_path = Path("./output/result.json").absolute()
        with open(config_path, "r", encoding="utf-8") as f:
            config_dict: dict = json.load(f)
        # 准备目录
        input_dir = Path("./input").absolute()
        output_dir = Path("./output").absolute()
        log_dir = Path("./log").absolute()
        input_dir.mkdir(parents=True, exist_ok=True)
        output_dir.mkdir(parents=True, exist_ok=True)
        log_dir.mkdir(parents=True, exist_ok=True)
        # 准备路径
        input_file = input_dir / config_dict.get("inputFile", "input.xlsx")
        input_csv = output_dir / "result.csv"
        output_pdf = output_dir / "result.pdf"
        output_png = output_dir / "result.png"
        work_sh = log_dir / "plot.sh"
        work_log = log_dir / "plot.log"

        # 转换输入文件
        convert_input(
            input_file,
            input_csv,
            p_value=config_dict.get("pValue", 0.05),
            fc=config_dict.get("fc", 2),
            top=config_dict["top"]
        )

        # 构建命令
        script_path = SCRIPTS / "Scatter" / "DiffRank" / "DiffRank.R"
        cmd = (
            f"#!/usr/bin/env bash\n\n"
            f"set -vex\n\n"
            f"{RSCRIPT} {script_path} "
            f"--config_json '{config_path}' "
            f"--input_csv '{input_csv}' "
            f"--output_pdf '{output_pdf}' "
            f"--output_png '{output_png}'\n\n"
            f"cd ./output \n\n"
            f"[ -f result.zip ] && rm result.zip \n\n"
            f"zip -r result.zip ./*\n\n"
        )

        # 执行
        click.echo(f"[✓] 开始执行绘图任务...")
        run_or_die(cmd=cmd, sh_file=work_sh, log_file=work_log)
        click.echo(f"[✓] 任务完成，结果已生成：\nPDF: {output_pdf}\nPNG: {output_png}")

    except Exception as e:
        click.echo(f"[✗] 绘图任务失败：{e}", err=True)
        raise click.Abort()
